import json
from pathlib import Path
import numpy as np
import torch
from torch import nn
from torch.nn import functional as F
from model_wrapper import FastRegressModelWrapper, SplitRegressModelWrapper


class Model(nn.Module):
	def __init__(self):
		# 调用父类的__init__()方法初始化, 这是必须的
		super().__init__()
		self.linear = nn.Linear(2, 2)
	
	def forward(self, x1: torch.Tensor, x2: torch.Tensor):
		return self.linear(x1 + x2)
	
	
if __name__ == '__main__':
	x1 = torch.randn(100, 2).tolist()
	x2 = torch.randn(100, 2).tolist()
	y = torch.randn(100, 2)
	
	# 构建model，训练
	model = Model()
	model_wrapper = SplitRegressModelWrapper(model)
	model_wrapper.train([x1, x2], y)
	